import os
import pandas as pd
import numpy as np
from scipy.stats import pearsonr

def calculate_metrics(precipitation, estimation):
    """
    计算平均偏差、平均绝对误差、均方根误差、平均相对误差、相关系数。

    参数:
    precipitation (pd.Series): 降水数据
    estimation (pd.Series): 估算数据

    返回:
    dict: 包含五个评估指标的字典
    """
    # 合并数据并删除包含缺失值、inf 和 -inf 的行
    combined = pd.concat([precipitation, estimation], axis=1)
    combined = combined.replace([np.inf, -np.inf], np.nan).dropna()
    precipitation = combined.iloc[:, 0]
    estimation = combined.iloc[:, 1]

    if len(precipitation) < 2 or len(estimation) < 2:
        # 数据长度不足 2，用 np.nan 填充指标
        return {
            '平均偏差': np.nan,
            '平均绝对误差': np.nan,
            '均方根误差': np.nan,
            '平均相对误差(%)': np.nan,
            '相关系数': np.nan
        }
    
    bias = np.mean(estimation - precipitation)
    mae = np.mean(np.abs(estimation - precipitation))
    rmse = np.sqrt(np.mean((estimation - precipitation) ** 2))
    
    # 过滤掉 precipitation 为 0 的数据
    non_zero_precipitation_mask = precipitation != 0
    non_zero_precipitation = precipitation[non_zero_precipitation_mask]
    non_zero_estimation = estimation[non_zero_precipitation_mask]
    
    if len(non_zero_precipitation) > 0:
        mre = np.mean(np.abs((non_zero_estimation - non_zero_precipitation) / non_zero_precipitation)) * 100
    else:
        mre = np.nan
    
    corr, _ = pearsonr(precipitation, estimation)
    return {
        '平均偏差': bias,
        '平均绝对误差': mae,
        '均方根误差': rmse,
        '平均相对误差(%)': mre,
        '相关系数': corr
    }




# 确保结果保存目录存在
result_dir = 'result'
output_dir = '统计结果/output'
os.makedirs(output_dir, exist_ok=True)





# 获取 result 目录下所有的 Excel 文件
excel_files = [f for f in os.listdir(result_dir) if f.endswith('.xlsx')]

for excel_file in excel_files:
    file_path = os.path.join(result_dir, excel_file)
    df = pd.read_excel(file_path)

    # 对“pre”和“IDW”进行统计分析
    # metrics_idw = calculate_metrics(df['pre'], df['IDW'])
    # 对“pre”和“Kriging”进行统计分析
    metrics_kriging = calculate_metrics(df['pre'], df['Kriging'])

    # 将结果转换为 DataFrame
    metrics_df = pd.DataFrame({
        # '指标': list(metrics_idw.keys()),
        # 'IDW 结果': list(metrics_idw.values()),
        'Kriging 结果': list(metrics_kriging.values())
    })

    # 生成输出文件名
    output_file_name = os.path.splitext(excel_file)[0] + '_统计结果.xlsx'
    output_file_path = os.path.join(output_dir, output_file_name)

    # 将结果保存到新的 Excel 文件
    metrics_df.to_excel(output_file_path, index=False)
    print(f"{excel_file} 的统计结果已保存到 {output_file_path}")